Skip to content

Conversation

skc7
Copy link
Contributor

@skc7 skc7 commented Jun 24, 2025

This PR introduces two new ops in omp dialect, omp.target_allocmem and omp.target_freemem.
omp.target_allocmem: Allocates heap memory on device. Will be lowered to omp_target_alloc call in llvm.
omp.target_freemem: Deallocates heap memory on device. Will be lowered to omp+target_free call in llvm.

Example:
%1 = omp.target_allocmem %device : i32, i64
omp.target_freemem %device, %1 : i32, i64

The work in this PR is C-P/inspired from @ivanradanov commit from coexecute implementation:
Add fir omp target alloc and free ops
Lower omp_target_{alloc,free} to llvm

This patch is pre-requisite for workdistribute lowering in #140523

@skc7 skc7 requested a review from mjklemm June 24, 2025 06:39
@skc7 skc7 requested a review from ivanradanov June 26, 2025 05:04
@skc7 skc7 self-assigned this Jun 26, 2025
@skc7 skc7 marked this pull request as ready for review June 26, 2025 08:20
@llvmbot llvmbot added flang Flang issues not falling into any other category flang:fir-hlfir flang:codegen labels Jun 26, 2025
@llvmbot
Copy link
Member

llvmbot commented Jun 26, 2025

@llvm/pr-subscribers-mlir-llvm
@llvm/pr-subscribers-mlir-openmp
@llvm/pr-subscribers-flang-fir-hlfir

@llvm/pr-subscribers-flang-codegen

Author: Chaitanya (skc7)

Changes

This PR introduces two new ops in fir dialect, fir.omp_target_allocmem and fir.omp_target_freemem.
fir.omp_target_allocmem: Allocates heap memory on device. Will be lowered to omp_target_alloc call in llvm.
fir.omp_target_freemem: Deallocates heap memory on device. Will be lowered to omp+target_free call in llvm.

Example:
%device = arith.constant 0 : i32
%1 = fir.omp_target_allocmem %device : i32, !fir.array<3x3xi32>
fir.omp_target_freemem %device, %1 : i32, !fir.heap<!fir.array<3x3xi32>>

The work in this PR is C-P/inspired from @ivanradanov commit from coexecute implementation:
Add fir omp target alloc and free ops
Lower omp_target_{alloc,free} to llvm


Full diff: https://github.com/llvm/llvm-project/pull/145464.diff

5 Files Affected:

  • (modified) flang/include/flang/Optimizer/Dialect/FIROps.td (+63)
  • (modified) flang/lib/Optimizer/CodeGen/CodeGen.cpp (+101-1)
  • (modified) flang/lib/Optimizer/Dialect/FIROps.cpp (+80-10)
  • (added) flang/test/Fir/omp_target_allocmem.fir (+28)
  • (added) flang/test/Fir/omp_target_freemem.fir (+28)
diff --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td
index 8ac847dd7dd0a..93d617027e30b 100644
--- a/flang/include/flang/Optimizer/Dialect/FIROps.td
+++ b/flang/include/flang/Optimizer/Dialect/FIROps.td
@@ -517,6 +517,69 @@ def fir_ZeroOp : fir_OneResultOp<"zero_bits", [NoMemoryEffect]> {
   let assemblyFormat = "type($intype) attr-dict";
 }
 
+def fir_OmpTargetAllocMemOp : fir_Op<"omp_target_allocmem",
+    [MemoryEffects<[MemAlloc<DefaultResource>]>, AttrSizedOperandSegments]> {
+  let summary = "allocate storage on an openmp device for an object of a given type";
+
+  let description = [{
+    Creates a heap memory reference suitable for storing a value of the
+    given type, T.  The heap refernce returned has type `!fir.heap<T>`.
+    The memory object is in an undefined state.  `omp_target_allocmem` operations must
+    be paired with `omp_target_freemem` operations to avoid memory leaks.
+
+    ```
+      %device = arith.constant 0 : i32
+      %1 = fir.omp_target_allocmem %device : i32, !fir.array<3x3xi32>
+    ```
+  }];
+
+  let arguments = (ins
+    Arg<AnyIntegerType>:$device,
+    TypeAttr:$in_type,
+    OptionalAttr<StrAttr>:$uniq_name,
+    OptionalAttr<StrAttr>:$bindc_name,
+    Variadic<AnyIntegerType>:$typeparams,
+    Variadic<AnyIntegerType>:$shape
+  );
+  let results = (outs fir_HeapType);
+
+  let hasCustomAssemblyFormat = 1;
+  let hasVerifier = 1;
+
+  let extraClassDeclaration = [{
+    mlir::Type getAllocatedType();
+    bool hasLenParams() { return !getTypeparams().empty(); }
+    bool hasShapeOperands() { return !getShape().empty(); }
+    unsigned numLenParams() { return getTypeparams().size(); }
+    operand_range getLenParams() { return getTypeparams(); }
+    unsigned numShapeOperands() { return getShape().size(); }
+    operand_range getShapeOperands() { return getShape(); }
+    static mlir::Type getRefTy(mlir::Type ty);
+  }];
+}
+
+def fir_OmpTargetFreeMemOp : fir_Op<"omp_target_freemem",
+  [MemoryEffects<[MemFree]>]> {
+  let summary = "free a heap object on an openmp device";
+
+  let description = [{
+    Deallocates a heap memory reference that was allocated by an `omp_target_allocmem`.
+    The memory object that is deallocated is placed in an undefined state
+    after `fir.omp_target_freemem`.
+    ```
+      %device = arith.constant 0 : i32
+      %1 = fir.omp_target_allocmem %device : i32, !fir.array<3x3xi32>
+      fir.omp_target_freemem %device, %1 : i32, !fir.heap<!fir.array<3x3xi32>>
+    ```
+  }];
+
+  let arguments = (ins
+  Arg<AnyIntegerType, "", [MemFree]>:$device,
+  Arg<fir_HeapType, "", [MemFree]>:$heapref
+  );
+  let assemblyFormat = "$device `,` $heapref attr-dict `:` type($device) `,` qualified(type($heapref))";
+}
+
 //===----------------------------------------------------------------------===//
 // Terminator operations
 //===----------------------------------------------------------------------===//
diff --git a/flang/lib/Optimizer/CodeGen/CodeGen.cpp b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
index a3de3ae9d116a..042ade6b1e0a1 100644
--- a/flang/lib/Optimizer/CodeGen/CodeGen.cpp
+++ b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
@@ -1168,6 +1168,105 @@ struct FreeMemOpConversion : public fir::FIROpConversion<fir::FreeMemOp> {
 };
 } // namespace
 
+static mlir::LLVM::LLVMFuncOp getOmpTargetAlloc(mlir::Operation *op) {
+  auto module = op->getParentOfType<mlir::ModuleOp>();
+  if (mlir::LLVM::LLVMFuncOp mallocFunc =
+          module.lookupSymbol<mlir::LLVM::LLVMFuncOp>("omp_target_alloc"))
+    return mallocFunc;
+  mlir::OpBuilder moduleBuilder(module.getBodyRegion());
+  auto i64Ty = mlir::IntegerType::get(module->getContext(), 64);
+  auto i32Ty = mlir::IntegerType::get(module->getContext(), 32);
+  return moduleBuilder.create<mlir::LLVM::LLVMFuncOp>(
+      moduleBuilder.getUnknownLoc(), "omp_target_alloc",
+      mlir::LLVM::LLVMFunctionType::get(
+          mlir::LLVM::LLVMPointerType::get(module->getContext()),
+          {i64Ty, i32Ty},
+          /*isVarArg=*/false));
+}
+
+namespace {
+struct OmpTargetAllocMemOpConversion
+    : public fir::FIROpConversion<fir::OmpTargetAllocMemOp> {
+  using FIROpConversion::FIROpConversion;
+
+  mlir::LogicalResult
+  matchAndRewrite(fir::OmpTargetAllocMemOp heap, OpAdaptor adaptor,
+                  mlir::ConversionPatternRewriter &rewriter) const override {
+    mlir::Type heapTy = heap.getType();
+    mlir::LLVM::LLVMFuncOp mallocFunc = getOmpTargetAlloc(heap);
+    mlir::Location loc = heap.getLoc();
+    auto ity = lowerTy().indexType();
+    mlir::Type dataTy = fir::unwrapRefType(heapTy);
+    mlir::Type llvmObjectTy = convertObjectType(dataTy);
+    if (fir::isRecordWithTypeParameters(fir::unwrapSequenceType(dataTy)))
+      TODO(loc, "fir.omp_target_allocmem codegen of derived type with length "
+                "parameters");
+    mlir::Value size = genTypeSizeInBytes(loc, ity, rewriter, llvmObjectTy);
+    if (auto scaleSize = genAllocationScaleSize(heap, ity, rewriter))
+      size = rewriter.create<mlir::LLVM::MulOp>(loc, ity, size, scaleSize);
+    for (mlir::Value opnd : adaptor.getOperands().drop_front())
+      size = rewriter.create<mlir::LLVM::MulOp>(
+          loc, ity, size, integerCast(loc, rewriter, ity, opnd));
+    auto mallocTyWidth = lowerTy().getIndexTypeBitwidth();
+    auto mallocTy =
+        mlir::IntegerType::get(rewriter.getContext(), mallocTyWidth);
+    if (mallocTyWidth != ity.getIntOrFloatBitWidth())
+      size = integerCast(loc, rewriter, mallocTy, size);
+    heap->setAttr("callee", mlir::SymbolRefAttr::get(mallocFunc));
+    rewriter.replaceOpWithNewOp<mlir::LLVM::CallOp>(
+        heap, ::getLlvmPtrType(heap.getContext()),
+        mlir::SmallVector<mlir::Value, 2>({size, heap.getDevice()}),
+        addLLVMOpBundleAttrs(rewriter, heap->getAttrs(), 2));
+    return mlir::success();
+  }
+
+  /// Compute the allocation size in bytes of the element type of
+  /// \p llTy pointer type. The result is returned as a value of \p idxTy
+  /// integer type.
+  mlir::Value genTypeSizeInBytes(mlir::Location loc, mlir::Type idxTy,
+                                 mlir::ConversionPatternRewriter &rewriter,
+                                 mlir::Type llTy) const {
+    return computeElementDistance(loc, llTy, idxTy, rewriter, getDataLayout());
+  }
+};
+} // namespace
+
+static mlir::LLVM::LLVMFuncOp getOmpTargetFree(mlir::Operation *op) {
+  auto module = op->getParentOfType<mlir::ModuleOp>();
+  if (mlir::LLVM::LLVMFuncOp freeFunc =
+          module.lookupSymbol<mlir::LLVM::LLVMFuncOp>("omp_target_free"))
+    return freeFunc;
+  mlir::OpBuilder moduleBuilder(module.getBodyRegion());
+  auto i32Ty = mlir::IntegerType::get(module->getContext(), 32);
+  return moduleBuilder.create<mlir::LLVM::LLVMFuncOp>(
+      moduleBuilder.getUnknownLoc(), "omp_target_free",
+      mlir::LLVM::LLVMFunctionType::get(
+          mlir::LLVM::LLVMVoidType::get(module->getContext()),
+          {getLlvmPtrType(module->getContext()), i32Ty},
+          /*isVarArg=*/false));
+}
+
+namespace {
+struct OmpTargetFreeMemOpConversion
+    : public fir::FIROpConversion<fir::OmpTargetFreeMemOp> {
+  using FIROpConversion::FIROpConversion;
+
+  mlir::LogicalResult
+  matchAndRewrite(fir::OmpTargetFreeMemOp freemem, OpAdaptor adaptor,
+                  mlir::ConversionPatternRewriter &rewriter) const override {
+    mlir::LLVM::LLVMFuncOp freeFunc = getOmpTargetFree(freemem);
+    mlir::Location loc = freemem.getLoc();
+    freemem->setAttr("callee", mlir::SymbolRefAttr::get(freeFunc));
+    rewriter.create<mlir::LLVM::CallOp>(
+        loc, mlir::TypeRange{},
+        mlir::ValueRange{adaptor.getHeapref(), freemem.getDevice()},
+        addLLVMOpBundleAttrs(rewriter, freemem->getAttrs(), 2));
+    rewriter.eraseOp(freemem);
+    return mlir::success();
+  }
+};
+} // namespace
+
 // Convert subcomponent array indices from column-major to row-major ordering.
 static llvm::SmallVector<mlir::Value>
 convertSubcomponentIndices(mlir::Location loc, mlir::Type eleTy,
@@ -4274,7 +4373,8 @@ void fir::populateFIRToLLVMConversionPatterns(
       GlobalLenOpConversion, GlobalOpConversion, InsertOnRangeOpConversion,
       IsPresentOpConversion, LenParamIndexOpConversion, LoadOpConversion,
       LocalitySpecifierOpConversion, MulcOpConversion, NegcOpConversion,
-      NoReassocOpConversion, SelectCaseOpConversion, SelectOpConversion,
+      NoReassocOpConversion, OmpTargetAllocMemOpConversion,
+      OmpTargetFreeMemOpConversion, SelectCaseOpConversion, SelectOpConversion,
       SelectRankOpConversion, SelectTypeOpConversion, ShapeOpConversion,
       ShapeShiftOpConversion, ShiftOpConversion, SliceOpConversion,
       StoreOpConversion, StringLitOpConversion, SubcOpConversion,
diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp
index ecfa2939e96a6..9335a4b041ac8 100644
--- a/flang/lib/Optimizer/Dialect/FIROps.cpp
+++ b/flang/lib/Optimizer/Dialect/FIROps.cpp
@@ -106,24 +106,38 @@ static bool verifyTypeParamCount(mlir::Type inType, unsigned numParams) {
   return false;
 }
 
-/// Parser shared by Alloca and Allocmem
-///
+/// Parser shared by Alloca, Allocmem and OmpTargetAllocmem
+/// boolean flag isTargetOp is used to identify omp_target_allocmem
 /// operation ::= %res = (`fir.alloca` | `fir.allocmem`) $in_type
 ///                      ( `(` $typeparams `)` )? ( `,` $shape )?
 ///                      attr-dict-without-keyword
+/// operation ::= %res = (`fir.omp_target_alloca`) $device : devicetype,
+///                      $in_type ( `(` $typeparams `)` )? ( `,` $shape )?
+///                      attr-dict-without-keyword
 template <typename FN>
-static mlir::ParseResult parseAllocatableOp(FN wrapResultType,
-                                            mlir::OpAsmParser &parser,
-                                            mlir::OperationState &result) {
+static mlir::ParseResult
+parseAllocatableOp(FN wrapResultType, mlir::OpAsmParser &parser,
+                   mlir::OperationState &result, bool isTargetOp = false) {
+  auto &builder = parser.getBuilder();
+  bool hasOperands = false;
+  std::int32_t typeparamsSize = 0;
+  // Parse device number as a new operand
+  if (isTargetOp) {
+    mlir::OpAsmParser::UnresolvedOperand deviceOperand;
+    mlir::Type deviceType;
+    if (parser.parseOperand(deviceOperand) || parser.parseColonType(deviceType))
+      return mlir::failure();
+    if (parser.resolveOperand(deviceOperand, deviceType, result.operands))
+      return mlir::failure();
+    if (parser.parseComma())
+      return mlir::failure();
+  }
   mlir::Type intype;
   if (parser.parseType(intype))
     return mlir::failure();
-  auto &builder = parser.getBuilder();
   result.addAttribute("in_type", mlir::TypeAttr::get(intype));
   llvm::SmallVector<mlir::OpAsmParser::UnresolvedOperand> operands;
   llvm::SmallVector<mlir::Type> typeVec;
-  bool hasOperands = false;
-  std::int32_t typeparamsSize = 0;
   if (!parser.parseOptionalLParen()) {
     // parse the LEN params of the derived type. (<params> : <types>)
     if (parser.parseOperandList(operands, mlir::OpAsmParser::Delimiter::None) ||
@@ -147,13 +161,19 @@ static mlir::ParseResult parseAllocatableOp(FN wrapResultType,
       parser.resolveOperands(operands, typeVec, parser.getNameLoc(),
                              result.operands))
     return mlir::failure();
+
   mlir::Type restype = wrapResultType(intype);
   if (!restype) {
     parser.emitError(parser.getNameLoc(), "invalid allocate type: ") << intype;
     return mlir::failure();
   }
-  result.addAttribute("operandSegmentSizes", builder.getDenseI32ArrayAttr(
-                                                 {typeparamsSize, shapeSize}));
+  llvm::SmallVector<std::int32_t> segmentSizes;
+  if (isTargetOp)
+    segmentSizes.push_back(1);
+  segmentSizes.push_back(typeparamsSize);
+  segmentSizes.push_back(shapeSize);
+  result.addAttribute("operandSegmentSizes",
+                      builder.getDenseI32ArrayAttr(segmentSizes));
   if (parser.parseOptionalAttrDict(result.attributes) ||
       parser.addTypeToList(restype, result.types))
     return mlir::failure();
@@ -385,6 +405,56 @@ llvm::LogicalResult fir::AllocMemOp::verify() {
   return mlir::success();
 }
 
+//===----------------------------------------------------------------------===//
+// OmpTargetAllocMemOp
+//===----------------------------------------------------------------------===//
+
+mlir::Type fir::OmpTargetAllocMemOp::getAllocatedType() {
+  return mlir::cast<fir::HeapType>(getType()).getEleTy();
+}
+
+mlir::Type fir::OmpTargetAllocMemOp::getRefTy(mlir::Type ty) {
+  return fir::HeapType::get(ty);
+}
+
+mlir::ParseResult
+fir::OmpTargetAllocMemOp::parse(mlir::OpAsmParser &parser,
+                                mlir::OperationState &result) {
+  return parseAllocatableOp(wrapAllocMemResultType, parser, result, true);
+}
+
+void fir::OmpTargetAllocMemOp::print(mlir::OpAsmPrinter &p) {
+  p << " ";
+  p.printOperand(getDevice());
+  p << " : ";
+  p << getDevice().getType();
+  p << ", ";
+  p << getInType();
+  if (!getTypeparams().empty()) {
+    p << '(' << getTypeparams() << " : " << getTypeparams().getTypes() << ')';
+  }
+  for (auto sh : getShape()) {
+    p << ", ";
+    p.printOperand(sh);
+  }
+  p.printOptionalAttrDict((*this)->getAttrs(),
+                          {"in_type", "operandSegmentSizes"});
+}
+
+llvm::LogicalResult fir::OmpTargetAllocMemOp::verify() {
+  llvm::SmallVector<llvm::StringRef> visited;
+  if (verifyInType(getInType(), visited, numShapeOperands()))
+    return emitOpError("invalid type for allocation");
+  if (verifyTypeParamCount(getInType(), numLenParams()))
+    return emitOpError("LEN params do not correspond to type");
+  mlir::Type outType = getType();
+  if (!mlir::dyn_cast<fir::HeapType>(outType))
+    return emitOpError("must be a !fir.heap type");
+  if (fir::isa_unknown_size_box(fir::dyn_cast_ptrEleTy(outType)))
+    return emitOpError("cannot allocate !fir.box of unknown rank or type");
+  return mlir::success();
+}
+
 //===----------------------------------------------------------------------===//
 // ArrayCoorOp
 //===----------------------------------------------------------------------===//
diff --git a/flang/test/Fir/omp_target_allocmem.fir b/flang/test/Fir/omp_target_allocmem.fir
new file mode 100644
index 0000000000000..5140c91c9510c
--- /dev/null
+++ b/flang/test/Fir/omp_target_allocmem.fir
@@ -0,0 +1,28 @@
+// RUN: %flang_fc1 -emit-llvm  %s -o - | FileCheck %s
+
+// CHECK-LABEL: define ptr @omp_target_allocmem_array_of_nonchar(
+// CHECK: call ptr @omp_target_alloc(i64 36, i32 0)
+func.func @omp_target_allocmem_array_of_nonchar() -> !fir.heap<!fir.array<3x3xi32>> {
+  %device = arith.constant 0 : i32
+  %1 = fir.omp_target_allocmem %device : i32, !fir.array<3x3xi32>
+  return %1 : !fir.heap<!fir.array<3x3xi32>>
+}
+
+// CHECK-LABEL: define ptr @omp_target_allocmem_array_of_char(
+// CHECK: call ptr @omp_target_alloc(i64 90, i32 0)
+func.func @omp_target_allocmem_array_of_char() -> !fir.heap<!fir.array<3x3x!fir.char<1,10>>> {
+  %device = arith.constant 0 : i32
+  %1 = fir.omp_target_allocmem %device : i32, !fir.array<3x3x!fir.char<1,10>>
+  return %1 : !fir.heap<!fir.array<3x3x!fir.char<1,10>>>
+}
+
+// CHECK-LABEL: define ptr @omp_target_allocmem_array_of_dynchar(
+// CHECK-SAME: i32 %[[len:.*]])
+// CHECK: %[[mul1:.*]] = sext i32 %[[len]] to i64
+// CHECK: %[[mul2:.*]] = mul i64 9, %[[mul1]]
+// CHECK: call ptr @omp_target_alloc(i64 %[[mul2]], i32 0)
+func.func @omp_target_allocmem_array_of_dynchar(%l: i32) -> !fir.heap<!fir.array<3x3x!fir.char<1,?>>> {
+  %device = arith.constant 0 : i32
+  %1 = fir.omp_target_allocmem %device : i32, !fir.array<3x3x!fir.char<1,?>>(%l : i32)
+  return %1 : !fir.heap<!fir.array<3x3x!fir.char<1,?>>>
+}
diff --git a/flang/test/Fir/omp_target_freemem.fir b/flang/test/Fir/omp_target_freemem.fir
new file mode 100644
index 0000000000000..02e136076a9cf
--- /dev/null
+++ b/flang/test/Fir/omp_target_freemem.fir
@@ -0,0 +1,28 @@
+// RUN: %flang_fc1 -emit-llvm  %s -o - | FileCheck %s
+
+// CHECK-LABEL: define void @omp_target_allocmem_array_of_nonchar(
+// CHECK: call void @omp_target_free(ptr {{.*}}, i32 0)
+func.func @omp_target_allocmem_array_of_nonchar() -> () {
+  %device = arith.constant 0 : i32
+  %1 = fir.omp_target_allocmem %device : i32, !fir.array<3x3xi32>
+  fir.omp_target_freemem %device, %1 : i32, !fir.heap<!fir.array<3x3xi32>>
+  return
+}
+
+// CHECK-LABEL: define void @omp_target_allocmem_array_of_char(
+// CHECK: call void @omp_target_free(ptr {{.*}}, i32 0)
+func.func @omp_target_allocmem_array_of_char() -> () {
+  %device = arith.constant 0 : i32
+  %1 = fir.omp_target_allocmem %device : i32, !fir.array<3x3x!fir.char<1,10>>
+  fir.omp_target_freemem %device, %1 : i32, !fir.heap<!fir.array<3x3x!fir.char<1,10>>>
+  return
+}
+
+// CHECK-LABEL: define void @omp_target_allocmem_array_of_dynchar(
+// CHECK: call void @omp_target_free(ptr {{.*}}, i32 0)
+func.func @omp_target_allocmem_array_of_dynchar(%l: i32) -> () {
+  %device = arith.constant 0 : i32
+  %1 = fir.omp_target_allocmem %device : i32, !fir.array<3x3x!fir.char<1,?>>(%l : i32)
+  fir.omp_target_freemem %device, %1 : i32, !fir.heap<!fir.array<3x3x!fir.char<1,?>>>
+  return
+}

Copy link
Contributor

@tblah tblah left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why did you decide to put this in the FIR dialect rather than OpenMP?

@ergawy
Copy link
Member

ergawy commented Jun 26, 2025

Thanks @skc7 for the PR! I have a few concerns about this approach, please let me know if I am missing anything:

  1. I am bit skeptical about leaking OpenMP specific ops to the fir dialect. The reason I see we have to do that is that we need to return a fir_HeapType from the omp_target_allocmem op, right? If there are not other reasons, why don't we introduce the op in the OpenMP dialect where it naturally fits, let it return an Int64 type and use fir.convert to the target type?
  2. In general though, why are we introducing a new op for the runtime API? Can't we directly emit these as call ops instead of having to introduce an op and then lower the op to a call? Are we planning to use these ops later somehow in some other way than lowering to the target call?

I am probably missing some context here though so my question might not be totally informed.

@skc7
Copy link
Contributor Author

skc7 commented Jun 27, 2025

@tblah @ergawy This PR is a pre-requisite for workdistribute construct implementation and lowering in flang.
#140523 introduces "lower-workdistribute" pass which in its implementation, moves fir.allocmem which is inside omp.target region to outside (to host). So, a new op which is called from host, but allocates memory on a given omp device is required.
So have added it in fir dialect following @ivanradanov coexecute implementation.

I would try to experiment adding this op to openmp mlir dialect as suggested and let you know if it fits this workdistribute implementation.

Copy link

github-actions bot commented Jul 3, 2025

✅ With the latest revision this PR passed the C/C++ code formatter.

@skc7
Copy link
Contributor Author

skc7 commented Jul 3, 2025

Thanks @skc7 for the PR! I have a few concerns about this approach, please let me know if I am missing anything:

  1. I am bit skeptical about leaking OpenMP specific ops to the fir dialect. The reason I see we have to do that is that we need to return a fir_HeapType from the omp_target_allocmem op, right? If there are not other reasons, why don't we introduce the op in the OpenMP dialect where it naturally fits, let it return an Int64 type and use fir.convert to the target type?
  2. In general though, why are we introducing a new op for the runtime API? Can't we directly emit these as call ops instead of having to introduce an op and then lower the op to a call? Are we planning to use these ops later somehow in some other way than lowering to the target call?

I am probably missing some context here though so my question might not be totally informed.

Hi @ergawy

For implementing workdistribute lowering, in a new pass "lower-workdistribute" : LINK, Need to move fir.allocmem nested in omp.target outside it (or hoist it outside omp.target)

Have updated the PR to implement the approach 1, to introduce omp.target_alloc_mem and omp.target_free_mem in openMP dialect. Also, their conversion to llvm ir has been implemented.
But, inorder to properly calculate the size of allocation of fir types, had to implement TargetAllocMemOpConversion in CodeGenOpenMP, to lower omp.target_alloc_mem to LLVM dialect callOp of "omp_target_alloc".

Approach 2, is to directly lower to a LLVM dialect, call to "omp_target_alloc".
Call to "omp_target_alloc" is to be made directly in the pass. It requires to calculate the actual size of allocation for fir type when it is lowered to llvm ir in the pass itself, which I think is bit complicated since it requires to move the utilities of size calculation (from Codegen fir to llvm) to with in the pass.

So, went ahead with approach 1. Please let me know your comments on this.

@skc7 skc7 changed the title [flang] Introduce omp_target_allocmem and omp_target_freemem fir ops. [flang] Introduce omp.target_allocmem and omp.target_freemem omp dialect ops. Jul 4, 2025
Copy link
Member

@ergawy ergawy left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the updates. I have some further comments/suggestions.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we are using a deprected conversion here (see: https://github.com/llvm/llvm-project/blob/main/llvm/include/llvm/Support/TypeSize.h#L360).

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It seems this function does more than what the name suggests (i.e. copies old attributes, drops the old operandSegmentSizes and recreates it, and adds and empty op_bundle_sizes).

Should we do the attribute changes in-place (i.e. at the call site) instead? Since this function does not seem reusable in other scenarios.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see this is adapted from CodeGen.cpp though, so see my comment below.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It seems this function and some of the ones below are shared between CodeGen.cpp and CodeGenOpenMP.cpp. Can we move them to a shared location, e.g. flang/Optimizer/Support/Utils.h/.cpp?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Moved these functions to utils.h in latest patch.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Windows build is failing due to linker error if I move this function to Utils. Linux build is failing in debug mode locally aswell. So, moved this back here as a static utility function to avoid any extra linking libs to FIRSupport.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
mlir::LLVM::LLVMPointerType::get(allocmemOp.getContext(), 0);
mlir::LLVM::LLVMPointerType::get(allocmemOp.getContext());

0 is the default value for the address space.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks. Updated.

Comment on lines +282 to +167
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this if condition is not covered by the introduced tests.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: move before use location below.

Copy link
Member

@ergawy ergawy Jul 9, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this conversion pattern should not replace omp.target_allocmem. Instead it should only convert the fir-specific allocation type to i64 and leave the lowering of omp.target_allocmem to call ptr @omp_target_alloc(...) to convertTargetAllocMemOp in OpenMPToLLVMIRTranslation.cpp. What we have now means that we do the same conversion in 2 different places. The preferred place would be OpenMPToLLVMIRTranslation.cpp since this is where all OpenMP to LLVM conversions are done.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @ergawy I agree that having same conversion in two different places is not ideal. But I'm having difficulty in figuring out, how would below flang fir type need to be converted to llvm types after the conversion?

omp.target_alloc_mem op takes a TypeAttr:$in_type (like fir.alloc_mem op), as argument. Now, in this conversion, fir array type needs to be replaced by i64 type as you have suggested. How would the resulting IR look like after conversion, for below IR?

func.func @omp_target_allocmem_array_of_dynchar(%l: i32) -> () {
%device = arith.constant 0 : i32
%1 = omp.target_allocmem %device : i32, !fir.array<3x3x!fir.char<1,?>>(%l : i32)
omp.target_freemem %device, %1 : i32, i64
return
}

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for looking into this Chaitanya. Here is a proposal on how we can do that:

diff --git a/flang/lib/Optimizer/CodeGen/CodeGenOpenMP.cpp b/flang/lib/Optimizer/CodeGen/CodeGenOpenMP.cpp
index 14cc7bb511f0..36f0265e8995 100644
--- a/flang/lib/Optimizer/CodeGen/CodeGenOpenMP.cpp
+++ b/flang/lib/Optimizer/CodeGen/CodeGenOpenMP.cpp
@@ -126,22 +126,6 @@ struct PrivateClauseOpConversion
   }
 };
 
-static mlir::LLVM::LLVMFuncOp getOmpTargetAlloc(mlir::Operation *op) {
-  auto module = op->getParentOfType<mlir::ModuleOp>();
-  if (mlir::LLVM::LLVMFuncOp mallocFunc =
-          module.lookupSymbol<mlir::LLVM::LLVMFuncOp>("omp_target_alloc"))
-    return mallocFunc;
-  mlir::OpBuilder moduleBuilder(module.getBodyRegion());
-  auto i64Ty = mlir::IntegerType::get(module->getContext(), 64);
-  auto i32Ty = mlir::IntegerType::get(module->getContext(), 32);
-  return moduleBuilder.create<mlir::LLVM::LLVMFuncOp>(
-      moduleBuilder.getUnknownLoc(), "omp_target_alloc",
-      mlir::LLVM::LLVMFunctionType::get(
-          mlir::LLVM::LLVMPointerType::get(module->getContext()),
-          {i64Ty, i32Ty},
-          /*isVarArg=*/false));
-}
-
 static mlir::Type convertObjectType(const fir::LLVMTypeConverter &converter,
                                     mlir::Type firType) {
   if (auto boxTy = mlir::dyn_cast<fir::BaseBoxType>(firType))
@@ -259,13 +243,11 @@ struct TargetAllocMemOpConversion
   matchAndRewrite(mlir::omp::TargetAllocMemOp allocmemOp, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     mlir::Type heapTy = allocmemOp.getAllocatedType();
-    mlir::LLVM::LLVMFuncOp mallocFunc = getOmpTargetAlloc(allocmemOp);
     mlir::Location loc = allocmemOp.getLoc();
     auto ity = lowerTy().indexType();
     mlir::Type dataTy = fir::unwrapRefType(heapTy);
     mlir::Type llvmObjectTy = convertObjectType(lowerTy(), dataTy);
-    mlir::Type llvmPtrTy =
-        mlir::LLVM::LLVMPointerType::get(allocmemOp.getContext(), 0);
+
     if (fir::isRecordWithTypeParameters(fir::unwrapSequenceType(dataTy)))
       TODO(loc, "omp.target_allocmem codegen of derived type with length "
                 "parameters");
@@ -273,6 +255,7 @@ struct TargetAllocMemOpConversion
                                           lowerTy().getDataLayout());
     if (auto scaleSize = genAllocationScaleSize(allocmemOp, ity, rewriter))
       size = rewriter.create<mlir::LLVM::MulOp>(loc, ity, size, scaleSize);
+
     for (mlir::Value opnd : adaptor.getOperands().drop_front())
       size = rewriter.create<mlir::LLVM::MulOp>(
           loc, ity, size, integerCast(lowerTy(), loc, rewriter, ity, opnd));
@@ -281,13 +264,13 @@ struct TargetAllocMemOpConversion
         mlir::IntegerType::get(rewriter.getContext(), mallocTyWidth);
     if (mallocTyWidth != ity.getIntOrFloatBitWidth())
       size = integerCast(lowerTy(), loc, rewriter, mallocTy, size);
-    allocmemOp->setAttr("callee", mlir::SymbolRefAttr::get(mallocFunc));
-    auto callOp = rewriter.create<mlir::LLVM::CallOp>(
-        loc, llvmPtrTy,
-        mlir::SmallVector<mlir::Value, 2>({size, allocmemOp.getDevice()}),
-        addLLVMOpBundleAttrs(rewriter, allocmemOp->getAttrs(), 2));
-    rewriter.replaceOpWithNewOp<mlir::LLVM::PtrToIntOp>(
-        allocmemOp, rewriter.getIntegerType(64), callOp.getResult());
+
+    rewriter.modifyOpInPlace(allocmemOp, [&]() {
+      allocmemOp.setInType(rewriter.getI8Type());
+      allocmemOp.getTypeparamsMutable().clear();
+      allocmemOp.getTypeparamsMutable().append(size);
+    });
+
     return mlir::success();
   }
 };
diff --git a/flang/test/Fir/omp_target_allocmem_freemem.fir b/flang/test/Fir/omp_target_allocmem_freemem.fir
index 920220272845..00a4462510a9 100644
--- a/flang/test/Fir/omp_target_allocmem_freemem.fir
+++ b/flang/test/Fir/omp_target_allocmem_freemem.fir
@@ -24,7 +24,8 @@ func.func @omp_target_allocmem_array_of_char() -> () {
 // CHECK-SAME: i32 %[[len:.*]])
 // CHECK: %[[mul1:.*]] = sext i32 %[[len]] to i64
 // CHECK: %[[mul2:.*]] = mul i64 9, %[[mul1]]
-// CHECK: call ptr @omp_target_alloc(i64 %[[mul2]], i32 0)
+// CHECK: %[[mul3:.*]] = mul i64 1, %[[mul2]]
+// CHECK: call ptr @omp_target_alloc(i64 %[[mul3]], i32 0)
 func.func @omp_target_allocmem_array_of_dynchar(%l: i32) -> () {
   %device = arith.constant 0 : i32
   %1 = omp.target_allocmem %device : i32, !fir.array<3x3x!fir.char<1,?>>(%l : i32)
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index c246d4abbdfe..0744bac98b3f 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -5923,7 +5923,12 @@ convertTargetAllocMemOp(Operation &opInst, llvm::IRBuilderBase &builder,
   mlir::Type heapTy = allocMemOp.getAllocatedType();
   llvm::Type *llvmHeapTy = moduleTranslation.convertType(heapTy);
   llvm::TypeSize typeSize = dataLayout.getTypeStoreSize(llvmHeapTy);
-  llvm::ConstantInt *allocSize = builder.getInt64(typeSize.getFixedValue());
+  llvm::Value *allocSize = builder.getInt64(typeSize.getFixedValue());
+
+  for (auto typeParam : allocMemOp.getTypeparams())
+    allocSize =
+        builder.CreateMul(allocSize, moduleTranslation.lookupValue(typeParam));
+
   // Create call to "omp_target_alloc" with the args as translated llvm values.
   llvm::CallInst *call =
       builder.CreateCall(ompTargetAllocFunc, {allocSize, llvmDeviceNum});

With that, after converting fir, we alway have a to omp.target_alloc_mem op that allocates the correct number of bytes based on the converted fir type.

So when TargetAllocMemOpConversion is run, you will get something like this:

  llvm.func @omp_target_allocmem_array_of_dynchar(%arg0: i32) {
    %0 = llvm.mlir.constant(0 : i32) : i32
    %1 = llvm.mlir.constant(1 : i64) : i64
    %2 = llvm.mlir.constant(9 : i64) : i64
    %3 = llvm.mul %1, %2 : i64
    %4 = llvm.sext %arg0 : i32 to i64
    %5 = llvm.mul %3, %4 : i64
    %6 = omp.target_allocmem %0 : i32, i8(%5 : i64)
    omp.target_freemem %0, %6 : i32, i64
    llvm.return
  }

(note that this is the conversion of the 3rd test you added in omp_target_allocmem_freemem.fir). You can reproduce the same result using: ./bin/fir-opt --cfg-conversion --fir-to-llvm-ir="target=aarch64-unknown-linux-gnu" which will test only this part of the pipeline. Then you can use mlir-translate to test the final OpenMP to LLVM conversion.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @ergawy. This does fix the issue I was facing. Updated the patch.

Copy link
Member

@ergawy ergawy left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @skc7 for taking care of the comments and apologies fir the delay. I just have a few more comments and this should be good to go.

However, it would be nice to get a review from some one else as well.

Comment on lines +1915 to +2165
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need both fields or is one sufficient? Also, please document the op fields in the description.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reference for this op is fir.allocmem and when fir.allocmem is hoisted out of omp.target region, it will be replaced by omp.target_allocmem in the workdistribute lowering pass in #140523

Have used the similar arguments for this op like in fir.allocmem.
$shape is required as it provide the runtime extents of dynamic dimensions when allocating Fortran arrays.

Updated the description in the latest patch.

Comment on lines 253 to 255
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you please de-templetize this and move the definiton to the .cpp file? From op, we only need InType and Loc; both of which can be easily passed to the function on call sites.

Adding templates in headers files (specially the ones who are used in a lot places like Uilts files) can slow down compilation times.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for feedback. Have updated this in latest patch.

@skc7 skc7 force-pushed the skc7/flang_workdistribute_preq_fir_targetops branch from 4314bfb to 6ef5914 Compare July 16, 2025 16:21
@skc7 skc7 requested a review from tblah July 22, 2025 12:54
@skc7 skc7 changed the title [flang] Introduce omp.target_allocmem and omp.target_freemem omp dialect ops. [OpenMP] Introduce omp.target_allocmem and omp.target_freemem omp dialect ops. Jul 22, 2025
@skc7 skc7 force-pushed the skc7/flang_workdistribute_preq_fir_targetops branch from 6ef5914 to 6c12425 Compare July 24, 2025 05:27
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These new omp ops needs to be translated first to llvm ir before its users gets translated. This is to avoid users querying these ops results in llvm ir before their translation itself.

@skc7 skc7 force-pushed the skc7/flang_workdistribute_preq_fir_targetops branch from e25b1ae to 5a2d87b Compare July 29, 2025 17:33
@skc7
Copy link
Contributor Author

skc7 commented Aug 4, 2025

Hi All,
This PR was initially intended to be for new fir ops but later changed to omp.target_allocmem and omp.target_freemem omp dialect ops as per feedback. Lowering of these ops to llvm ir has also been implemented.
Please review the PR and suggest any changes that are required to be done.
Thanks.

@mjklemm mjklemm requested a review from ergawy August 4, 2025 11:48
Copy link
Member

@ergawy ergawy left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM! Thanks for the updates.

However, it would be nice to have another approval (even if only a high-level review).

@kparzysz kparzysz requested a review from skatrak August 4, 2025 12:41
Copy link
Contributor

@mjklemm mjklemm left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@skc7 skc7 merged commit 4a3bf27 into llvm:main Aug 18, 2025
9 checks passed
searlmc1 pushed a commit to ROCm/llvm-project that referenced this pull request Aug 18, 2025
… omp dialect ops. (llvm#145464)"

needs downstream PR to land first

This reverts commit 4a3bf27.
searlmc1 pushed a commit to ROCm/llvm-project that referenced this pull request Aug 19, 2025
…lect ops. (llvm#145464)

This PR introduces two new ops in omp dialect, omp.target_allocmem and
omp.target_freemem.
omp.target_allocmem: Allocates heap memory on device. Will be lowered to
omp_target_alloc call in llvm.
omp.target_freemem: Deallocates heap memory on device. Will be lowered
to omp+target_free call in llvm.

Example:
  %1 = omp.target_allocmem %device : i32, i64
  omp.target_freemem %device, %1 : i32, i64

The work in this PR is C-P/inspired from @ivanradanov commit from
coexecute implementation:
[Add fir omp target alloc and free
ops](ivanradanov@be860ac)
[Lower omp_target_{alloc,free} to
llvm](ivanradanov@6e2d584)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants